#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-08-29

import os
import threading
import time
import cv2
import uuid
import os
import platform
import argparse
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, QoS, BaseNode, get_extra_args, Rate
from pathlib import Path
import json


class ClassificationDatasetReaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True,
            **kwargs
        )
        self.dataset_dir = self.get_param("dataset_dir", "/home/jario/dataset/ILSVRC2012_img_val")
        self.class_codes_map = self.get_param("class_codes_map", "imagenet_class_codes_map.json")
        self.class_ids_map = self.get_param("class_ids_map", "imagenet_class_ids_map.json")
        self.auto_next = self.get_param("auto_next", False)
        self.repeated = self.get_param("repeated", False)
        self.rate = self.get_param("rate", 2)
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.queue_pool.append(self.img_i_queue)
        self.img_i = 0
        self.frame_id = 0
        self.done = False

        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        if len(self.class_codes_map) > 0:
            if not self.class_codes_map.startswith("/"):
                self.class_codes_map = os.path.join(params_dir, self.class_codes_map)
            with open(self.class_codes_map, 'r', encoding='utf-8') as f:
                self.class_codes_dict = json.load(f)

        if len(self.class_ids_map) > 0:
            if not self.class_ids_map.startswith("/"):
                self.class_ids_map = os.path.join(params_dir, self.class_ids_map)
            with open(self.class_ids_map, 'r', encoding='utf-8') as f:
                self.class_ids_dict = json.load(f)

        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        
        class_directories = os.listdir(self.dataset_dir)
        class_directories = sorted(class_directories)
        self.img_fns = []
        for class_name in class_directories:
            p = Path(os.path.join(self.dataset_dir, class_name))
            if p.is_dir():
                img_directories = os.listdir(p)
                img_directories = sorted(img_directories)
                for img_name in img_directories:
                    if os.path.splitext(img_name)[-1] in ['.jpg', '.jpeg', '.JPEG', '.png']:
                        img_fn = os.path.join(self.dataset_dir, class_name, img_name)
                        if len(self.class_codes_map) > 0:
                            mapped_class_name = self.class_codes_dict[class_name]
                        else:
                            mapped_class_name = class_name
                        if len(self.class_ids_map) > 0:
                            class_id = self.class_ids_dict[mapped_class_name]
                        else:
                            class_id = int(mapped_class_name)
                        self.img_fns.append([img_fn, class_id, mapped_class_name])

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._next_reader.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.img_fns):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1
                if self.repeated and self.img_i == len(self.img_fns):
                    self.img_i = 0
            else:
                self.done = True

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img = cv2.imread(self.img_fns[img_i][0])

            if self.b_use_shm:
                msg = self._image_writer.cvimg2sms_mem(img)
            else:
                msg = cvimg2sms(img)

            msg['img_id'] = img_i
            msg['img_total'] = len(self.img_fns)
            msg['file_name'] = os.path.basename(self.img_fns[img_i][0])
            msg['client_id'] = self.client_id
            msg['class_id'] = self.img_fns[img_i][1]
            msg['class_name'] = self.img_fns[img_i][2]
            self._image_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = ClassificationDatasetReaderNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.launch_next()
    if node.auto_next:
        r = Rate(node.rate)
        while not node.done:
            node.launch_next()
            r.sleep()
        node.shutdown()
    node.spin()
